package yycore

import chisel3._
import chisel3.util._
import chisel3.util.experimental.BoringUtils
import common.Constants._
import bus.CoreLinkIO
import utils.LookupTree



class WBCommitBus extends Bundle() {
  val instValid = Output(Bool())
  val inst = Output(UInt(32.W))
  val pc = Output(UInt(AddrBits.W))
  val rf_wen = Output(Bool())
  val wb_addr = Output(UInt(5.W))
  val wb_data = Output(UInt(XLEN.W))
  val is_mmio = Output(Bool())
  val intrNO = Output(UInt(XLEN.W))
  val redirect = new RedirectIO
}

class EXEStage() extends Module {
  val io = IO(new Bundle() {
    val flush = Input(Bool())
    val in = Flipped(Decoupled(new DecodeCtrlSignal))
    val out = Decoupled(new WBCommitBus())
    val bypass = Flipped(new ByPassDataIO)
    val dmem = new CoreLinkIO(XLEN)
    val mmio = new CoreLinkIO(XLEN)
  })
  io := DontCare

  val s_exe_ready = Wire(Bool())
  io.in.ready := s_exe_ready

  //**********************************
  // Pipeline State Registers
  // Execute State
//  val reset_exe = WireInit(0.U.asTypeOf(new DecodeCtrlSignal))
//  reset_exe.inst := BUBBLE
//  val exe_reg = RegInit(reset_exe)
//
//  //**********************************
//  // update State Registers
//  when(io.out.fire() && io.in.fire() && !io.flush){
//    exe_reg := io.in.bits
//  }.elsewhen(io.flush || (io.out.fire() && !io.in.fire())){
//    exe_reg := reset_exe
//  }
  val in = io.in

  val (src1, src2) = (in.bits.op1_data, in.bits.op2_data)
  // **********************************
  // exe function
  val fun = in.bits.fu_fun
  val fu_type = in.bits.fu_typ
  val isALU = FuType.isALU(fu_type) && !io.flush//&& exe_reg.instValid
  val isLSU = FuType.isLSU(fu_type) && !io.flush//&& exe_reg.instValid
  val isCSR = FuType.isCSR(fu_type) && !io.flush//&& exe_reg.instValid
  val isMOU = FuType.isMOU(fu_type) && !io.flush//&& exe_reg.instValid
  //val isMDU = FuType.isMDU(fu_type) && !io.flush//&& exe_reg.instValid

  val lsOccur = RegInit(false.B)


  /**
   * ALU
   */
  val alu = Module(new ALU())
  val aluOut = alu.access(isALU, src1, src2, fun)
  alu.io.out.ready := io.out.ready

  /**
   * MDU
   */
//    val mdu = Module(new MDU)
//    val mduOut = mdu.access(isMDU, src1, src2, fun)
//    mdu.io.out.ready := io.out.ready

  /**
   * CSR
   */
  val csr = Module(new CSR())
  csr.io.req.valid := isCSR
  csr.io.req.bits.addr := Mux(fun === CSROpType.N && isCSR, 0.U, in.bits.inst(31, 20))
  csr.io.req.bits.cmd := Mux(isCSR, fun, CSROpType.N) //exe_reg.csr_cmd
  csr.io.req.bits.wdata := src1
  csr.io.req.bits.rd := in.bits.wb_addr
  csr.io.req.bits.pc := in.bits.pc
  csr.io.req.bits.instValid := in.bits.instValid && !lsOccur && io.out.fire() && !io.flush
  csr.io.req.bits.exceptionVec := in.bits.exceptionVec
  val csr_rdata = csr.io.resp.bits.rdata

  /**
   * LSU
   */
  val lsu = Module(new LSU())
  val lsuOut = lsu.access(valid = isLSU && !csr.io.raiseIntr, (src1 + src2).asUInt(), fun, in.bits.rs2_data)
  lsu.io.out.ready := io.out.ready
  lsu.io.dmem <> io.dmem
  lsu.io.mmio <> io.mmio
  when(lsu.io.in.fire()) { lsOccur := true.B }
  when(lsu.io.out.fire()) { lsOccur := false.B }

  /**
   * MOU
   */
  val mou = Module(new MOU)
  mou.io.req.valid := isMOU
  mou.io.req.bits.fun := fun
  mou.io.req.bits.pc := in.bits.pc

  // Output
  val exe_out = Mux(isALU, aluOut, lsuOut)//Mux(isLSU, lsuOut, mduOut))
  val outValid = MuxCase(true.B, Array(
    (isLSU && !csr.io.raiseIntr)  -> lsu.io.out.valid
  ))

  s_exe_ready := io.out.fire()

  val br_eq = (in.bits.op1_data === in.bits.rs2_data)
  val br_lt = (in.bits.op1_data.asSInt() < in.bits.rs2_data.asSInt())
  val br_ltu = (in.bits.op1_data.asUInt() < in.bits.rs2_data.asUInt())

  // ***********************************************
  // Branch Logic
  val exe_ctrl_pc_sel = MuxCase(PC_4, Array(
    (in.bits.br_type === BRType.none) ->  PC_4,
    (in.bits.br_type === BRType.bne)  ->  Mux(!br_eq, PC_BR, PC_4),
    (in.bits.br_type === BRType.beq)  ->  Mux(br_eq, PC_BR, PC_4),
    (in.bits.br_type === BRType.bge)  ->  Mux(!br_lt, PC_BR, PC_4),
    (in.bits.br_type === BRType.bgeu) ->  Mux(!br_ltu, PC_BR, PC_4),
    (in.bits.br_type === BRType.blt)  ->  Mux(br_lt, PC_BR, PC_4),
    (in.bits.br_type === BRType.bltu) ->  Mux(br_ltu, PC_BR, PC_4),
    (in.bits.br_type === BRType.jal)  ->  PC_J,
    (in.bits.br_type === BRType.jalr) ->  PC_JR,
    (in.bits.br_type === BRType.call) ->  PC_J,
    (in.bits.br_type === BRType.ret)  ->  PC_JR,
  ))

  def taken(pc_sel: UInt) = pc_sel =/= PC_4

  val target = Mux(BRType.isBType(in.bits.br_type), in.bits.pc + in.bits.op2_data, exe_out)
  val actualTarget = Mux(taken(exe_ctrl_pc_sel), target, in.bits.pc + 4.U)
  // val actualTarget = MuxCase(0.U, Array(
  //   (exe_ctrl_pc_sel === PC_BR)   -> exe_br_jmp_target,
  //   (exe_ctrl_pc_sel === PC_J)    -> exe_jump_reg_target,
  //   (exe_ctrl_pc_sel === PC_JR)   -> exe_jump_reg_target,
  // ))
  val preTarget = in.bits.prepcNext
  val predictWrong = BRType.isBranch(in.bits.br_type) && preTarget =/= actualTarget
  
  val bpuUpdateReq = WireInit(0.U.asTypeOf(new BPUUpdateReq))
  bpuUpdateReq.valid := in.bits.instValid && in.bits.br_type =/= BR_N && io.out.fire() && !io.flush
  bpuUpdateReq.pc := in.bits.pc
  bpuUpdateReq.isMissPredict := predictWrong
  bpuUpdateReq.actualTarget := actualTarget
  bpuUpdateReq.actualTaken := taken(exe_ctrl_pc_sel)
  bpuUpdateReq.btbType := LookupTree(in.bits.br_type, BRType.bruFuncTobtbTypeTable)
  bpuUpdateReq.brType := in.bits.br_type
  BoringUtils.addSource(bpuUpdateReq, "bpuUpdateReq")

  // ******************
  // Branch and Jump Target Calculation
  io.out.bits.redirect.valid := predictWrong || csr.io.redirect.valid || mou.io.redirect.valid
  io.out.bits.redirect.target := Mux(mou.io.redirect.valid, mou.io.redirect.target,
    Mux(csr.io.redirect.valid, csr.io.redirect.target, actualTarget))


  // ****************************
  // output next stage singals
  val wb_data = MuxCase(aluOut, Array(
    (in.bits.wb_sel === WB_EXE)  ->  exe_out,
    (in.bits.wb_sel === WB_PC4)  ->  (in.bits.pc + 4.U),
    (in.bits.wb_sel === WB_CSR)  ->  csr_rdata
  ))
  val next_rf_wen = in.bits.rf_wen && !csr.io.wenFix
  // ===============================================
  // bypass output
  io.bypass.apply(rf_wen = next_rf_wen, addr = in.bits.wb_addr, data = wb_data)

  io.out.valid := outValid
  io.out.bits.instValid := in.bits.instValid
  io.out.bits.inst := in.bits.inst
  io.out.bits.pc := in.bits.pc
  io.out.bits.wb_addr := in.bits.wb_addr
  io.out.bits.rf_wen := next_rf_wen
  io.out.bits.wb_data := wb_data
  io.out.bits.is_mmio := lsu.io.isMMIO
  io.out.bits.intrNO := csr.io.resp.bits.intrNO

}
